package org.example.weather.bigscreen.core.component;

import com.google.common.base.Preconditions;
import lombok.extern.slf4j.Slf4j;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.logging.jdbc.PreparedStatementLogger;
import org.apache.ibatis.plugin.*;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.SqlSessionFactory;
import org.apache.shardingsphere.driver.jdbc.core.connection.ShardingSphereConnection;
import org.apache.shardingsphere.driver.jdbc.core.statement.ShardingSpherePreparedStatement;
import org.apache.shardingsphere.infra.binder.LogicSQL;
import org.apache.shardingsphere.infra.binder.SQLStatementContextFactory;
import org.apache.shardingsphere.infra.binder.aware.ParameterAware;
import org.apache.shardingsphere.infra.binder.segment.insert.keygen.GeneratedKeyContext;
import org.apache.shardingsphere.infra.binder.statement.SQLStatementContext;
import org.apache.shardingsphere.infra.binder.statement.dml.InsertStatementContext;
import org.apache.shardingsphere.infra.context.kernel.KernelProcessor;
import org.apache.shardingsphere.infra.database.type.DatabaseTypeEngine;
import org.apache.shardingsphere.infra.datanode.DataNode;
import org.apache.shardingsphere.infra.executor.check.SQLCheckEngine;
import org.apache.shardingsphere.infra.executor.sql.context.ExecutionContext;
import org.apache.shardingsphere.infra.parser.ShardingSphereSQLParserEngine;
import org.apache.shardingsphere.mode.metadata.MetaDataContexts;
import org.apache.shardingsphere.parser.rule.SQLParserRule;
import org.apache.shardingsphere.sql.parser.sql.common.statement.SQLStatement;
import org.redisson.api.RedissonClient;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Lazy;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.stereotype.Component;

import java.lang.reflect.Field;
import java.sql.Statement;
import java.util.*;
import java.util.concurrent.CompletionService;
import java.util.concurrent.ExecutorCompletionService;
import java.util.concurrent.Future;
import java.util.stream.Collectors;

@Slf4j
@Component
@Intercepts({@Signature(type = StatementHandler.class, method = "update", args = {Statement.class}), @Signature(type = StatementHandler.class, method = "query", args = {Statement.class, ResultHandler.class})})
public class ShardingRouteContextInterceptor implements Interceptor {
    @Value("${shardingsphere.datasource.tables.cache-key-prefix:%s_%s_cache_set}")
    private String dataSourceTablesCacheKey;
    private final CompletionService<Collection<DataNode>> completionService;
    private final Map<String, SqlSessionFactory> sqlSessionFactoryMap;
    private final RedissonClient redissonClient;

    public ShardingRouteContextInterceptor(@Lazy @Qualifier("sqlSessionFactoryMap") final Map<String, SqlSessionFactory> sqlSessionFactoryMap, final ThreadPoolTaskExecutor threadPoolTaskExecutor, final RedissonClient redissonClient) {
        this.sqlSessionFactoryMap = sqlSessionFactoryMap;
        this.completionService = new ExecutorCompletionService<>(threadPoolTaskExecutor);
        this.redissonClient = redissonClient;
    }

    @Override
    public Object intercept(final Invocation invocation) throws Throwable {
        Object[] args = invocation.getArgs();
        Object target = invocation.getTarget();
        if (target instanceof StatementHandler) {
            StatementHandler statementHandler = (StatementHandler) target;
            Field field = args[0].getClass().getSuperclass().getDeclaredField("h");
            field.setAccessible(true);
            Object obj = field.get(args[0]);
            if (obj instanceof PreparedStatementLogger) {
                PreparedStatementLogger preparedStatementLogger = (PreparedStatementLogger) obj;
                if(!(preparedStatementLogger.getPreparedStatement() instanceof ShardingSpherePreparedStatement)) {
                    return invocation.proceed();
                }
                ShardingSpherePreparedStatement shardingSpherePreparedStatement = (ShardingSpherePreparedStatement) preparedStatementLogger.getPreparedStatement();
                ShardingSphereConnection connection = shardingSpherePreparedStatement.getConnection();
                MetaDataContexts metaDataContexts = connection.getContextManager().getMetaDataContexts();
                Optional<SQLParserRule> sqlParserRule = metaDataContexts.getMetaData().getGlobalRuleMetaData().findSingleRule(SQLParserRule.class);
                Preconditions.checkState(sqlParserRule.isPresent());
                ShardingSphereSQLParserEngine sqlParserEngine = sqlParserRule.get().getSQLParserEngine(
                        DatabaseTypeEngine.getTrunkDatabaseTypeName(metaDataContexts.getMetaData().getDatabases().get(connection.getDatabaseName()).getResource().getDatabaseType()));
                SQLStatement sqlStatement = sqlParserEngine.parse(statementHandler.getBoundSql().getSql(), true);
                SQLStatementContext<?> sqlStatementContext = SQLStatementContextFactory.newInstance(metaDataContexts.getMetaData().getDatabases(), sqlStatement, connection.getDatabaseName());
                KernelProcessor kernelProcessor = new KernelProcessor();
                LogicSQL logicSQL = createLogicSQL(shardingSpherePreparedStatement.getParameters(), sqlStatementContext, statementHandler.getBoundSql().getSql());
                Collection<Comparable<?>> generatedValues = new LinkedList<>();
                ExecutionContext executionContext = createExecutionContext(logicSQL, metaDataContexts, kernelProcessor, connection, generatedValues);
                Map<String, List<DataNode>> todoMap = executionContext.getRouteContext().getOriginalDataNodes().stream().flatMap(Collection::stream).collect(Collectors.groupingBy(DataNode::getDataSourceName));
                todoMap.forEach((k, v) -> {
                    completionService.submit(new AsyncCreateTableTask(v, sqlSessionFactoryMap.get(k), redissonClient, dataSourceTablesCacheKey, k));
                });
                for (int i = 0; i < todoMap.size(); i++) {
                    try {
                        Future<Collection<DataNode>> result = completionService.take();
                        Collection<DataNode> dataNodes = result.get();
                        for (final DataNode dataNode : dataNodes) {
                            log.debug("--------------------------------------------------------------------------");
                            log.debug(dataNode.getDataSourceName());
                            log.debug(dataNode.getSchemaName());
                            log.debug(dataNode.getTableName());
                            log.debug("--------------------------------------------------------------------------");
                        }
                    } catch (InterruptedException e) {
                        log.error("completionService.take faild", e);
                    }
                }
            }
        }
        return invocation.proceed();
    }

    @Override
    public Object plugin(final Object target) {
        // 只对要拦截制定类型的对象生成代理
        if (target instanceof StatementHandler) {
            // 调用插件
            return Plugin.wrap(target, this);
        }
        return target;
    }

    @Override
    public void setProperties(final Properties properties) {

    }

    private LogicSQL createLogicSQL(List<Object> origionalParameters, SQLStatementContext<?> sqlStatementContext, String sql) {
        List<Object> parameters = new ArrayList<>(origionalParameters);
        if (sqlStatementContext instanceof ParameterAware) {
            ((ParameterAware) sqlStatementContext).setUpParameters(parameters);
        }
        return new LogicSQL(sqlStatementContext, sql, parameters);
    }

    private ExecutionContext createExecutionContext(final LogicSQL logicSQL, final MetaDataContexts metaDataContexts, final KernelProcessor kernelProcessor, final ShardingSphereConnection connection, final Collection<Comparable<?>> generatedValues) {
        SQLCheckEngine.check(logicSQL.getSqlStatementContext().getSqlStatement(), logicSQL.getParameters(),
                             metaDataContexts.getMetaData().getDatabases().get(connection.getDatabaseName()).getRuleMetaData().getRules(),
                             connection.getDatabaseName(), metaDataContexts.getMetaData().getDatabases(), null);
        ExecutionContext result = kernelProcessor.generateExecutionContext(
                logicSQL, metaDataContexts.getMetaData().getDatabases().get(connection.getDatabaseName()), metaDataContexts.getMetaData().getProps());
        findGeneratedKey(result).ifPresent(generatedKey -> generatedValues.addAll(generatedKey.getGeneratedValues()));
        return result;
    }

    private Optional<GeneratedKeyContext> findGeneratedKey(final ExecutionContext executionContext) {
        return executionContext.getSqlStatementContext() instanceof InsertStatementContext
                ? ((InsertStatementContext) executionContext.getSqlStatementContext()).getGeneratedKeyContext()
                : Optional.empty();
    }
}
